import pandas as pd
import json
import os
import cv2
import random
from dr_to_jpg import monochrome_convert,read_dcm_from_dir
from tqdm import tqdm
Point_Name_Dict={
    "{'关键点': ['右侧肺尖']}":0,
    "{'关键点': ['左侧肺尖']}":1,
    "{'关键点': ['右侧侧肋膈角']}":2,
    "{'关键点': ['左侧侧肋膈角']}":3,
    "{'关键点': ['第一胸椎上缘中点']}":4,
    "{'关键点': ['第二胸椎上缘中点']}":5,
    "{'关键点': ['第一胸椎上缘右侧']}":6,
    "{'关键点': ['第一胸椎上缘左侧']}":7,
    "{'关键点': ['第二胸椎上缘右侧']}":8,
    "{'关键点': ['第二胸椎上缘左侧']}":9,
}


def conver_csv(csv_dir, save_json_dir, jpg_path):
    ## 将csv文件进行拆分。，分成三个任务（关键点，检测，关键的连接而成的mask区域）分开保存到json中。
    ## 并将dcm文件转换成jpg进行保存。
    point_dict = {}
    csv_data = pd.read_csv(csv_dir)
    csv_data_point = csv_data[(csv_data['geometry'] == 'point')]
    csv_data_point = csv_data_point[(csv_data_point['comment'] == "['']")]
    ## 处理point
    last_name = csv_data_point.iloc[0].loc['sop']
    point_dict[csv_data_point.iloc[0].loc['sop']] = []
    for index in range(len(csv_data_point)):
        if last_name != csv_data_point.iloc[index].loc['sop']:
            point_dict[csv_data_point.iloc[index].loc['sop']] = []
        point = [int(csv_data_point.iloc[index].loc['points'].split(',')[0].replace('[[', '')),
                 int(csv_data_point.iloc[index].loc['points'].split(',')[1].replace(']]', '').replace(' ', ''))]
        point_dict[csv_data_point.iloc[index].loc['sop']].append(
            [point, Point_Name_Dict[csv_data_point.iloc[index].loc['selected']]])
        last_name = csv_data_point.iloc[index].loc['sop']
    
    json_list = []
    for k, value in point_dict.items():
        # if os.path.exists(os.path.join(jpg_path, k + '.jpg')):
            for i in value:
                if i[0][0] < 0 or i[0][1] < 0:
                    value.remove(i)
            ## 针对存在缺失值的数据 进行处理.
            value_name = [p[1] for p in value]
            true_name = [i for i in range(10)]
            miss_name = [p for p in true_name if p not in value_name]
            for j in miss_name:
                value.append([[0, 0], j])
            json_list.append({
                'name': k,
                'point': sorted(value, key=lambda j: j[1])
            })
    random.shuffle(json_list)
    print(len(json_list))
    with open(save_json_dir, 'w') as f:
        json.dump(json_list, f)
    with open(save_json_dir.replace('.json', '_train.json'), 'w') as f:
        json.dump(json_list[:int(0.8 * len(json_list))], f)
    with open(save_json_dir.replace('.json', '_val.json'), 'w') as f:
        json.dump(json_list[int(0.2 * len(json_list)):], f)



def show_point(image_path,json_dir):
    json_data = json.load(open(json_dir))
    for anno in json_data:
        points = [p[0] for p in anno['point']]
        try:
            image = cv2.imread(os.path.join(image_path,anno['name']+'.jpg'))
            for point in points:
                cv2.circle(image, (point[0],point[1]), 5, (255, 128, 255), thickness=3)
            cv2.imshow('image',cv2.resize(image,(512,512)))
            cv2.waitKey()
        except Exception as e:
            print(e)
            
def save_pic(dcm_path,save_path):
    os.makedirs(save_path,exist_ok=True)
    for path, dir_list, file_list in tqdm(os.walk(dcm_path)):
        for file_name in file_list:
            try:
                dir = os.path.join(path, file_name)
                image,spacing = read_dcm_from_dir(dir)
                print(spacing)
                image = monochrome_convert(image)
                cv2.imwrite(os.path.join(save_path,file_name.split('_')[1].replace('dcm','jpg')),image)
            except Exception as e:
                print(e)
                continue

if __name__ == '__main__':
    # save_pic(dcm_path='/home/lmy/PycharmProjects/Chest_point/data_0304/',save_path='data_0304/images')
    conver_csv(csv_dir='data_0304/Chest_CT_localizer_20210303.csv',save_json_dir='label_file_0304/point_0304.json',jpg_path=None)
    # show_point(image_path='data_0304/images',json_dir='label_file_0304/point_0304.json')